package com.lst.selfaiagent.chatMemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.KryoException;
import com.lst.selfaiagent.utils.TimestampedMessage;
import com.lst.selfaiagent.utils.TimestampedMessageConversation;
import org.objenesis.strategy.InstantiatorStrategy;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.PromptTemplate;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 基于文件持久化的对话记忆
 */
public class FileBasedChatMemory implements ChatMemory {

    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy( new StdInstantiatorStrategy());
        // 注册 TimestampedMessage 类以提高序列化效率
        kryo.register(TimestampedMessage.class);
        kryo.register(LocalDateTime.class);
        kryo.register(TimestampedMessageConversation.class); // 新增注册
        kryo.register(ArrayList.class); // 注册 ArrayList

    }

    // 构造对象时，指定文件保存目录
    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }

    // 文件存储
    @Override
    public void add(String conversationId, List<Message> messages) {
//        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        List<TimestampedMessage> conversationMessages = getOrCreateConversation(conversationId);
        // 为每个消息添加时间戳
        List<TimestampedMessage> timestampedMessages = messages.stream()
                .map(TimestampedMessage::new)
                .collect(Collectors.toList());
        conversationMessages.addAll(timestampedMessages);
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId) {
        List<TimestampedMessage> allMessages = getOrCreateConversation(conversationId);
        return allMessages.stream()
                .map(TimestampedMessage::getMessage) // 提取原始消息
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }


    private List<TimestampedMessage> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<TimestampedMessage> messages = new ArrayList<>();
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                TimestampedMessageConversation conversation = kryo.readObject(input, TimestampedMessageConversation.class);
                return conversation.getMessages();
            } catch (KryoException e) {
                // 处理Kryo相关的异常，如BufferUnderflowException
                System.err.println("Kryo序列化错误，可能文件已损坏，将使用空的对话历史: " + e.getMessage());
                // 文件可能已损坏，删除它以便重新开始
                file.delete();
            } catch (IOException e) {
                System.err.println("IO错误: " + e.getMessage());
            }
        }
        return messages;
    }

    private void saveConversation(String conversationId, List<TimestampedMessage> messages) {
        File file = getConversationFile(conversationId);
        TimestampedMessageConversation timestampedMessageConversation = new TimestampedMessageConversation();
        timestampedMessageConversation.setMessages(messages);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, timestampedMessageConversation);
        } catch (IOException e) {
            System.err.println("保存对话时发生错误: " + e.getMessage());
        }
    }

    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }

    // 新增方法：获取带时间戳的消息
    public List<TimestampedMessage> getTimestampedMessages(String conversationId, int lastN) {
        List<TimestampedMessage> allMessages = getOrCreateConversation(conversationId);
        return allMessages.stream()
                .skip(Math.max(0, allMessages.size() - lastN))
                .collect(Collectors.toList());
    }
}